// Copyright (c) 2024 Huawei Technologies Co., Ltd.
// openFuyao is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//          http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
// EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
// MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.

package lm

import (
	corev1 "k8s.io/api/core/v1"
)

const (
	nfdLabelPrefix       = "feature.node.kubernetes.io/"
	npuDeviceClassId     = "1200" // Processing accelerators
	npuDeviceVendorId    = "19e5"
	npuDeviceLabelPrefix = nfdLabelPrefix + "pci-" + npuDeviceClassId + "_" + npuDeviceVendorId + "_"
	npuDeviceLabelValue  = "true"
	npuPresentLabelKey   = "openfuyao.com/npu.present"
)

var (
	npuDeviceTypes = map[string]string{
		npuDeviceLabelPrefix + "d100.present": "huawei-Ascend310",
		npuDeviceLabelPrefix + "d500.present": "huawei-Ascend310P",
		npuDeviceLabelPrefix + "d801.present": "huawei-Ascend910",
		npuDeviceLabelPrefix + "d802.present": "huawei-Ascend910",
	}
)

func newAscendPresenceLabeler(node *corev1.Node) (Labeler, bool) {
	hasNPU := hasNPUDeviceLabels(node.Labels)

	if hasNPU {
		return Labels{npuPresentLabelKey: ""}, hasNPU
	} else {
		return nil, hasNPU
	}
}

func hasNPUDeviceLabels(labels map[string]string) bool {
	for key, value := range labels {
		_, ok := npuDeviceTypes[key]
		if ok && value == npuDeviceLabelValue {
			return true
		}
	}
	return false
}
